progress_bars.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280
  1. from __future__ import division
  2. import itertools
  3. import sys
  4. from signal import SIGINT, default_int_handler, signal
  5. from pip._vendor import six
  6. from pip._vendor.progress.bar import Bar, FillingCirclesBar, IncrementalBar
  7. from pip._vendor.progress.spinner import Spinner
  8. from pip._internal.utils.compat import WINDOWS
  9. from pip._internal.utils.logging import get_indentation
  10. from pip._internal.utils.misc import format_size
  11. from pip._internal.utils.typing import MYPY_CHECK_RUNNING
  12. if MYPY_CHECK_RUNNING:
  13. from typing import Any, Dict, List
  14. try:
  15. from pip._vendor import colorama
  16. # Lots of different errors can come from this, including SystemError and
  17. # ImportError.
  18. except Exception:
  19. colorama = None
  20. def _select_progress_class(preferred, fallback):
  21. # type: (Bar, Bar) -> Bar
  22. encoding = getattr(preferred.file, "encoding", None)
  23. # If we don't know what encoding this file is in, then we'll just assume
  24. # that it doesn't support unicode and use the ASCII bar.
  25. if not encoding:
  26. return fallback
  27. # Collect all of the possible characters we want to use with the preferred
  28. # bar.
  29. characters = [
  30. getattr(preferred, "empty_fill", six.text_type()),
  31. getattr(preferred, "fill", six.text_type()),
  32. ]
  33. characters += list(getattr(preferred, "phases", []))
  34. # Try to decode the characters we're using for the bar using the encoding
  35. # of the given file, if this works then we'll assume that we can use the
  36. # fancier bar and if not we'll fall back to the plaintext bar.
  37. try:
  38. six.text_type().join(characters).encode(encoding)
  39. except UnicodeEncodeError:
  40. return fallback
  41. else:
  42. return preferred
  43. _BaseBar = _select_progress_class(IncrementalBar, Bar) # type: Any
  44. class InterruptibleMixin(object):
  45. """
  46. Helper to ensure that self.finish() gets called on keyboard interrupt.
  47. This allows downloads to be interrupted without leaving temporary state
  48. (like hidden cursors) behind.
  49. This class is similar to the progress library's existing SigIntMixin
  50. helper, but as of version 1.2, that helper has the following problems:
  51. 1. It calls sys.exit().
  52. 2. It discards the existing SIGINT handler completely.
  53. 3. It leaves its own handler in place even after an uninterrupted finish,
  54. which will have unexpected delayed effects if the user triggers an
  55. unrelated keyboard interrupt some time after a progress-displaying
  56. download has already completed, for example.
  57. """
  58. def __init__(self, *args, **kwargs):
  59. # type: (List[Any], Dict[Any, Any]) -> None
  60. """
  61. Save the original SIGINT handler for later.
  62. """
  63. # https://github.com/python/mypy/issues/5887
  64. super(InterruptibleMixin, self).__init__( # type: ignore
  65. *args,
  66. **kwargs
  67. )
  68. self.original_handler = signal(SIGINT, self.handle_sigint)
  69. # If signal() returns None, the previous handler was not installed from
  70. # Python, and we cannot restore it. This probably should not happen,
  71. # but if it does, we must restore something sensible instead, at least.
  72. # The least bad option should be Python's default SIGINT handler, which
  73. # just raises KeyboardInterrupt.
  74. if self.original_handler is None:
  75. self.original_handler = default_int_handler
  76. def finish(self):
  77. # type: () -> None
  78. """
  79. Restore the original SIGINT handler after finishing.
  80. This should happen regardless of whether the progress display finishes
  81. normally, or gets interrupted.
  82. """
  83. super(InterruptibleMixin, self).finish() # type: ignore
  84. signal(SIGINT, self.original_handler)
  85. def handle_sigint(self, signum, frame): # type: ignore
  86. """
  87. Call self.finish() before delegating to the original SIGINT handler.
  88. This handler should only be in place while the progress display is
  89. active.
  90. """
  91. self.finish()
  92. self.original_handler(signum, frame)
  93. class SilentBar(Bar):
  94. def update(self):
  95. # type: () -> None
  96. pass
  97. class BlueEmojiBar(IncrementalBar):
  98. suffix = "%(percent)d%%"
  99. bar_prefix = " "
  100. bar_suffix = " "
  101. phases = (u"\U0001F539", u"\U0001F537", u"\U0001F535") # type: Any
  102. class DownloadProgressMixin(object):
  103. def __init__(self, *args, **kwargs):
  104. # type: (List[Any], Dict[Any, Any]) -> None
  105. # https://github.com/python/mypy/issues/5887
  106. super(DownloadProgressMixin, self).__init__( # type: ignore
  107. *args,
  108. **kwargs
  109. )
  110. self.message = (" " * (
  111. get_indentation() + 2
  112. )) + self.message # type: str
  113. @property
  114. def downloaded(self):
  115. # type: () -> str
  116. return format_size(self.index) # type: ignore
  117. @property
  118. def download_speed(self):
  119. # type: () -> str
  120. # Avoid zero division errors...
  121. if self.avg == 0.0: # type: ignore
  122. return "..."
  123. return format_size(1 / self.avg) + "/s" # type: ignore
  124. @property
  125. def pretty_eta(self):
  126. # type: () -> str
  127. if self.eta: # type: ignore
  128. return "eta {}".format(self.eta_td) # type: ignore
  129. return ""
  130. def iter(self, it): # type: ignore
  131. for x in it:
  132. yield x
  133. # B305 is incorrectly raised here
  134. # https://github.com/PyCQA/flake8-bugbear/issues/59
  135. self.next(len(x)) # noqa: B305
  136. self.finish()
  137. class WindowsMixin(object):
  138. def __init__(self, *args, **kwargs):
  139. # type: (List[Any], Dict[Any, Any]) -> None
  140. # The Windows terminal does not support the hide/show cursor ANSI codes
  141. # even with colorama. So we'll ensure that hide_cursor is False on
  142. # Windows.
  143. # This call needs to go before the super() call, so that hide_cursor
  144. # is set in time. The base progress bar class writes the "hide cursor"
  145. # code to the terminal in its init, so if we don't set this soon
  146. # enough, we get a "hide" with no corresponding "show"...
  147. if WINDOWS and self.hide_cursor: # type: ignore
  148. self.hide_cursor = False
  149. # https://github.com/python/mypy/issues/5887
  150. super(WindowsMixin, self).__init__(*args, **kwargs) # type: ignore
  151. # Check if we are running on Windows and we have the colorama module,
  152. # if we do then wrap our file with it.
  153. if WINDOWS and colorama:
  154. self.file = colorama.AnsiToWin32(self.file) # type: ignore
  155. # The progress code expects to be able to call self.file.isatty()
  156. # but the colorama.AnsiToWin32() object doesn't have that, so we'll
  157. # add it.
  158. self.file.isatty = lambda: self.file.wrapped.isatty()
  159. # The progress code expects to be able to call self.file.flush()
  160. # but the colorama.AnsiToWin32() object doesn't have that, so we'll
  161. # add it.
  162. self.file.flush = lambda: self.file.wrapped.flush()
  163. class BaseDownloadProgressBar(WindowsMixin, InterruptibleMixin,
  164. DownloadProgressMixin):
  165. file = sys.stdout
  166. message = "%(percent)d%%"
  167. suffix = "%(downloaded)s %(download_speed)s %(pretty_eta)s"
  168. class DefaultDownloadProgressBar(BaseDownloadProgressBar,
  169. _BaseBar):
  170. pass
  171. class DownloadSilentBar(BaseDownloadProgressBar, SilentBar):
  172. pass
  173. class DownloadBar(BaseDownloadProgressBar,
  174. Bar):
  175. pass
  176. class DownloadFillingCirclesBar(BaseDownloadProgressBar,
  177. FillingCirclesBar):
  178. pass
  179. class DownloadBlueEmojiProgressBar(BaseDownloadProgressBar,
  180. BlueEmojiBar):
  181. pass
  182. class DownloadProgressSpinner(WindowsMixin, InterruptibleMixin,
  183. DownloadProgressMixin, Spinner):
  184. file = sys.stdout
  185. suffix = "%(downloaded)s %(download_speed)s"
  186. def next_phase(self):
  187. # type: () -> str
  188. if not hasattr(self, "_phaser"):
  189. self._phaser = itertools.cycle(self.phases)
  190. return next(self._phaser)
  191. def update(self):
  192. # type: () -> None
  193. message = self.message % self
  194. phase = self.next_phase()
  195. suffix = self.suffix % self
  196. line = ''.join([
  197. message,
  198. " " if message else "",
  199. phase,
  200. " " if suffix else "",
  201. suffix,
  202. ])
  203. self.writeln(line)
  204. BAR_TYPES = {
  205. "off": (DownloadSilentBar, DownloadSilentBar),
  206. "on": (DefaultDownloadProgressBar, DownloadProgressSpinner),
  207. "ascii": (DownloadBar, DownloadProgressSpinner),
  208. "pretty": (DownloadFillingCirclesBar, DownloadProgressSpinner),
  209. "emoji": (DownloadBlueEmojiProgressBar, DownloadProgressSpinner)
  210. }
  211. def DownloadProgressProvider(progress_bar, max=None): # type: ignore
  212. if max is None or max == 0:
  213. return BAR_TYPES[progress_bar][1]().iter
  214. else:
  215. return BAR_TYPES[progress_bar][0](max=max).iter